{ "cells": [ { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "## Neural physical engines for inferring the halo mass distribution function\n", "\n", "Tom Charnock\n", "\n", "Institut d'Astrophysique de Paris\n", "


\n", "Notebook: presentations.charnock.fr/NPE\n", "\n", "\n", " \n", "
\n", " \"Sorbonne\n", " \n", " \"ANR\"\n", " \n", " \"IAP\"\n", " \n", " \"CNRS\"\n", " \n", " \"Aquila\"\n", "
" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "## Cosmology from the universe\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "#### Bayesian Origins Reconstruction of Galaxies (BORG)\n", "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "#### Infer the dark matter field from any catalogue\n", "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Halo mass distribution function" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Distribution of halos\n", "- Press Schechter (gravitational collapse following power law at small masses, exponential cut-off)\n", "- Local density deviations\n", "- Unknown stochasticity (Poissionian)" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Halo mass distribution\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Neural bias model\n", "






\n", "github:neural_bias_model" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Neural physical engines\n", "\n", "Build networks using physical principles.\n", "\n", "- Reduces number of parameters\n", "\n", "- Increases computational efficiency\n", "\n", "- Decreases overfitting\n", "\n", "- Improves interpretability" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Convolutions for translational invariance\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Architecture suited to the data\n", "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Use only causally relevant features\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Use only informative features\n", "
\n", "\n", "
\n", "github:multipole_kernels" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# What's the best map from the dark matter density today to some summary of the halos?" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### We don't actually know yet..." ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "#### We do know that it should be:\n", "- Dependent on the local patch\n", "- Mostly sensitive to the amplitude\n", "- Independent of position on the sky" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## Our neural physical engine\n", "
\n", "\n", "
\n", "$${\\Tiny \\psi_j^{\\ell,m} = A\\left(\\sum_{i=-\\kappa/2}^{i=\\kappa/2}K_{i}^{\\ell,m}\\delta_{j-i}^\\textrm{LPT}+b_j\\right), ~~~ \\ell=0, m=0}$$" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Neural density estimators" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "#### Halo mass distribution function is a smooth function of mass given a density environment" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "# Use a mixture density network" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Our neural density estimator\n", "$$\\begin{align*}\n", " {\\tiny n(M|\\delta) =}&{\\tiny \\sum_i^N\\alpha(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\alpha}})\\mathcal{N}\\left(\\mu(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\mu}}), \\sigma(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\sigma}})|M\\right),}\\\\\n", " {\\tiny =}&{\\tiny \\sum_i^N\\frac{\\alpha(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\alpha}})}{\\sqrt{2\\pi\\left(\\sigma(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\sigma}})\\right)^2}}\\exp\\left[-\\frac{\\left(\\log(M) - \\mu(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\mu}})\\right)^2}{2\\left(\\sigma(\\boldsymbol{\\psi}, \\boldsymbol{\\theta}_i^{\\boldsymbol{\\sigma}})\\right)^2}\\right],}\\end{align*}\n", "$$\n", "
\n", " \n", "
\n", "
\n", " $$\n", " \\begin{align*}\n", "{\\tiny\\alpha_i} & {\\tiny=\\text{softplus}(w_{i}^\\alpha\\boldsymbol{\\psi}+b^\\alpha_i),}\\\\\n", "{\\tiny\\mu_i }&{\\tiny = \\left\\{\\begin{array}{ll}\n", " w_{i}^\\mu\\boldsymbol{\\psi}+b^\\mu_i&i=0\\\\\n", " \\textrm{Max}\\left[0,~ w_{i}^\\mu\\boldsymbol{\\psi}+b^\\mu_i\\right]+\\mu_{i-1}&i>0\\\\\n", " \\end{array}\\right.,}\\\\\n", "{\\tiny\\sigma_i} &{\\tiny = \\text{softplus}(w_{i}^\\sigma\\boldsymbol{\\psi}+b^\\sigma_i)\\;,}\n", " \\end{align*}\n", " $$\n", "
\n", "
\n", "2 Gaussians with
${\\Tiny\\boldsymbol{b}^\\alpha \\to \\boldsymbol{b}^\\alpha +\\log(10^{-3})}$, ${\\Tiny b_0^\\mu \\to b_0^\\mu + \\log\\left(2\\times10^{12}\\right)}$ and ${\\Tiny\\boldsymbol{b}^\\sigma\\to \\boldsymbol{b}^\\alpha + \\log(10^3)}$.\n", "
" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Likelihood" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "$$\\begin{align*}\n", " {\\Tiny\\mathcal{L} =}&{\\Tiny \\sum_{j\\in\\textrm{catalogue}}\\log\\left[\\sum_i^N\\frac{\\alpha_{i,j}}{\\sqrt{2\\pi\\sigma_{i,j}^2}}\\exp\\left[-\\frac{\\left(\\log(M_j) - \\mu_{i,j}\\right)^2}{2\\sigma_{i,j}^2}\\right]\\right]}\\\\\n", " &{\\Tiny - V\\sum_{j\\in\\textrm{voxels},i=1}^N\\frac{\\alpha_{i,j}}{2}\\exp\\left[\\frac{\\sigma_{i,j}^2}{2}\\right]\\textrm{erfc}\\left[\\frac{\\log\\left(M_\\textrm{th}\\right) - \\mu_{i,j} - \\sigma_{i,j}^2}{\\sqrt{2\\sigma_{i,j}^2}}\\right].}\n", " \\end{align*}$$" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "This is Poissonian.
\n", "But on the NPE field, so it is non-local" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# HMCLET" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### We want to infer the weights of the neural bias model" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "### Hamiltonian Monte Carlo\n", "
\n", "Introduce momentum ${\\small {\\bf p}\\leftarrow\\mathcal{N}({\\bf 0},{\\bf M})}$ and solve Hamilton's equations\n", "

\n", "$$\\begin{align*}\n", "{\\Tiny\\mathcal{H}(\\boldsymbol{\\theta}, {\\bf p}) }&{\\Tiny~= \\mathcal{V}(\\boldsymbol{\\theta}) + \\mathcal{K}({\\bf p})}\\\\\n", "&{\\Tiny~=\\mathcal{L}(\\boldsymbol{\\theta}|\\boldsymbol{\\delta})-\\log\\left[\\pi(\\boldsymbol{\\theta})\\right]+\\frac{1}{2}{\\bf p}^T{\\bf M}^{-1}{\\bf p}}\n", "\\end{align*}$$\n", "Accept samples according to probability\n", "\n", "$$\\begin{equation*}\n", " {\\Tiny\\alpha = \\textrm{Min}\\left[\\exp\\left(\\Delta\\mathcal{H}\\right), 1\\right]}\n", "\\end{equation*}$$\n", "
\n", "
\n", "

\n", " \n", "
\n", "
\n", "
\n", "Evolve using ${\\Tiny\\dot{\\boldsymbol{\\theta}}= {\\bf M}^{-1}{\\bf p}}$ and ${\\Tiny\\dot{\\bf p}= -\\nabla\\mathcal{V}(\\boldsymbol{\\theta})}$\n", "
" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Acceptance is guaranteed if ODEs are perfectly solved" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "# Use leapfrog algorithm" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Leapfrog algorithm\n", "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Acceptance criterion\n", "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# How does this work for neural networks?" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## Likelihood surface is extremely flat and highly degenerate\n", "
\n", "\n", "\n", "Nearly impossible to know the mass matrix, ${\\bf M}$." ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## Use second order geometric information\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Quasi-Newtonian inverse Hessian\n", "\n", "$$\\begin{equation*}\n", " {\\Tiny{\\bf B}^*=\\left(\\mathbb{I}-\\frac{\\boldsymbol{s}_i\\boldsymbol{y}_i^T}{\\boldsymbol{y}_i^T\\boldsymbol{s}_i}\\right){\\bf B}_i\\left(\\mathbb{I}-\\frac{\\boldsymbol{y}_i\\boldsymbol{s}_i^T}{\\boldsymbol{y}_i^T\\boldsymbol{s}_i}\\right)+\\frac{\\boldsymbol{s}_i\\boldsymbol{s}_i^T}{\\boldsymbol{s}_i^T\\boldsymbol{y}_i},}\n", " \\end{equation*}\n", "$$\n", "with ${\\Tiny\\boldsymbol{s}_i=\\boldsymbol{\\theta}^*-\\boldsymbol{\\theta}_i}$ and ${\\Tiny\\boldsymbol{y}_i=\\nabla\\mathcal{L}(\\boldsymbol{\\theta}^*|\\boldsymbol{\\delta})-\\nabla\\mathcal{L}(\\boldsymbol{\\theta}_i|\\boldsymbol{\\delta})}$.\n", "\n", "#### Both $\\boldsymbol{s}$ and $\\boldsymbol{y}$ come for free with the leapfrog algorithm" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## We can now successfully traverse the likelihood surface\n", "\n", "The momenta and trajectories get rescaled by the Hessian providing a surface which is not as highly degenerate or flat" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Actual results" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "## VELMASS simulation\n", "\n", "- 125h-1Mpc chunk from 2h-1Gpc length box
\n", "- Rebinned to 643 voxels
\n", "- Planck-like cosmology
\n", "- Halos found using Rockstar ~ O(104) halos\n", "- Minimum halo size of 8x108 solar masses" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Run BORG" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "\n", "" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Sample the power spectrum\n", "
\n", "\n", "✓ Power spectra remain smooth and close to prior" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "#### Sample the weights of the neural bias model\n", "
\n", "\n", "\n", "✓ Weight values are properly sampled after burn in
\n", "✓ NPE acts as a contrast enhancer" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### See the effects of the neural physical engine\n", "
\n", "\n", "\n", "✓ Non-local information is used to improve fit" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# And what does the halo mass distribution function look like?" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "#### Halo mass distribution function from neural bias model\n", "
\n", "\n", "✓ Fits data!" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "# Conclusions" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### First example of zero-shot training, fully Bayesian neural network" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "fragment" } }, "source": [ "- Using neural physical engines to keep number of parameters small\n", "- Using the HMCLET to sample the weights of the neural network even though the likelihood is flat and degenerate\n", "- Using BORG to infer the simulations" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "subslide" } }, "source": [ "### Future work\n", "- Skip the bias completely by flowing the dark matter distribution to the stochastically sampled data (inferring the likelihood)\n", "- Optimisation of the neural physical engine architecture for optimal halo mass distribution function" ] }, { "cell_type": "markdown", "metadata": { "slideshow": { "slide_type": "slide" } }, "source": [ "## Neural physical engines for inferring the halo mass distribution function\n", "\n", "Tom Charnock\n", "\n", "Institut d'Astrophysique de Paris\n", "


\n", "\n", "\n", "\n", " \n", "
\n", " \"Sorbonne\n", " \n", " \"ANR\"\n", " \n", " \"IAP\"\n", " \n", " \"CNRS\"\n", " \n", " \"Aquila\"\n", "
" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "celltoolbar": "Diaporama", "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.7.9" } }, "nbformat": 4, "nbformat_minor": 2 }